# import cv2
# import numpy as np
# import os

# def crop_black_border(image_path, output_path):
#     """ 去除图片四周的黑色边框 """
#     img = cv2.imread(image_path)
#     if img is None:
#         print(f"无法读取图像：{image_path}")
#         return

#     # 转换为灰度图
#     gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

#     # 二值化，确保黑色背景为 0，其他部分为 255
#     _, binary = cv2.threshold(gray, 10, 255, cv2.THRESH_BINARY)

#     # 查找轮廓
#     contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

#     if not contours:
#         print(f"未找到有效内容，跳过：{image_path}")
#         return

#     # 找到最大轮廓区域
#     largest_contour = max(contours, key=cv2.contourArea)
#     x, y, w, h = cv2.boundingRect(largest_contour)

#     # 裁剪图像
#     cropped = img[y:y+h, x:x+w]

#     # 保存裁剪后的图像
#     cv2.imwrite(output_path, cropped)
#     print(f"已裁剪：{image_path} → {output_path}")

# # 批量处理数据集
# def batch_process_images(input_dir, output_dir):
#     """ 批量处理文件夹中的所有图片 """
#     if not os.path.exists(output_dir):
#         os.makedirs(output_dir)

#     for filename in os.listdir(input_dir):
#         if filename.lower().endswith(('.png', '.jpg', '.jpeg')):
#             input_path = os.path.join(input_dir, filename)
#             output_path = os.path.join(output_dir, filename)
#             crop_black_border(input_path, output_path)

# # 运行
# input_directory = "/mnt/e/pycharmworkspace/Race/waibao/FundAi-Bitirme-Projesi-main/datasets/Training_Dataset"  # 替换为你的输入文件夹
# output_directory = "/mnt/e/pycharmworkspace/Race/waibao/FundAi-Bitirme-Projesi-main/datasets/proposed"  # 替换为你的输出文件夹
# batch_process_images(input_directory, output_directory)
import cv2
import numpy as np
import os
import random

def horizontal_flip(image):
    """水平翻转"""
    return cv2.flip(image, 1)

def vertical_flip(image):
    """垂直翻转"""
    return cv2.flip(image, 0)

def random_rotate(image, angle):
    """旋转指定角度"""
    (h, w) = image.shape[:2]
    center = (w // 2, h // 2)
    M = cv2.getRotationMatrix2D(center, angle, 1.0)
    rotated = cv2.warpAffine(image, M, (w, h))
    return rotated

def augment_image(image, output_dir, base_name, ext, category, num_needed):
    """对图像进行数据增强，直到满足 3000 张"""
    augmentations = [
        ("hflip", horizontal_flip(image)),
        ("vflip", vertical_flip(image)),
        ("rot_15", random_rotate(image, 15)),
        ("rot_30", random_rotate(image, 30)),
        ("rot_45", random_rotate(image, 45)),
        ("rot_60", random_rotate(image, 60)),
        ("rot_90", random_rotate(image, 90)),
        ("rot_180", random_rotate(image, 180))
    ]

    saved_count = 0  # 计数已经保存的增强图片数量
    while saved_count < num_needed:
        for aug_name, aug_img in augmentations:
            if saved_count >= num_needed:
                break  # 需要的数量已满足
            aug_filename = f"{base_name}_{aug_name}{ext}"
            save_path = os.path.join(output_dir, aug_filename)
            cv2.imwrite(save_path, aug_img)
            saved_count += 1

def process_category(category_path, output_category_path, max_images=3000):
    """处理单个类别文件夹，增强图片直到 3000 张"""
    if not os.path.exists(output_category_path):
        os.makedirs(output_category_path)

    # 获取当前类别的所有图片
    image_files = [f for f in os.listdir(category_path) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
    num_original = len(image_files)

    print(f"类别 {os.path.basename(category_path)}: 原始图片 {num_original} 张")

    # 直接拷贝原始图片到目标文件夹
    for filename in image_files:
        src_path = os.path.join(category_path, filename)
        dst_path = os.path.join(output_category_path, filename)
        image = cv2.imread(src_path)
        if image is None:
            print(f"无法读取图像: {src_path}")
            continue
        cv2.imwrite(dst_path, image)

    # 如果原始图片已满足 max_images，则直接返回
    if num_original >= max_images:
        return

    # **需要填充的图片数量**
    num_needed = max_images - num_original
    print(f"类别 {os.path.basename(category_path)} 需要增加 {num_needed} 张")

    # 计算每张原始图片需要变换多少次
    num_images = len(image_files)
    num_transforms_per_image = (max_images - num_images) // num_images + 1

    # **数据增强**
    for img_name in image_files:
        base_name, ext = os.path.splitext(img_name)
        img_path = os.path.join(category_path, img_name)

        image = cv2.imread(img_path)
        if image is None:
            print(f"无法读取图像: {img_path}")
            continue

        augment_image(image, output_category_path, base_name, ext, category_path, num_transforms_per_image)

    print(f"类别 {os.path.basename(category_path)} 处理完成，共 {max_images} 张图片")

def batch_augment_images(input_dir, output_dir, max_images_per_class=3000):
    """批量处理所有类别，确保每个类别至少有 max_images_per_class 张"""
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    for category in os.listdir(input_dir):
        category_path = os.path.join(input_dir, category)

        # 只处理目录
        if os.path.isdir(category_path):
            output_category_path = os.path.join(output_dir, category)
            process_category(category_path, output_category_path, max_images_per_class)

# 运行
input_directory = "/mnt/e/pycharmworkspace/Race/waibao/FundAi-Bitirme-Projesi-main/datasets/proposed"  # 输入文件夹（有8个类别子文件夹）
output_directory = "/mnt/e/pycharmworkspace/Race/waibao/FundAi-Bitirme-Projesi-main/datasets/proposed_Aug"  # 输出文件夹
batch_augment_images(input_directory, output_directory)
